I am thinking about the following API change for this PR:

    // API on `master`
    DEPRECATED(LLAMA_API void llama_kv_cache_clear(ctx));
    DEPRECATED(LLAMA_API bool llama_kv_cache_seq_rm(ctx));
    DEPRECATED(LLAMA_API void llama_kv_cache_seq_cp(ctx));
    DEPRECATED(LLAMA_API void llama_kv_cache_seq_keep(ctx));
    DEPRECATED(LLAMA_API void llama_kv_cache_seq_add(ctx));
    DEPRECATED(LLAMA_API void llama_kv_cache_seq_div(ctx));
    DEPRECATED(LLAMA_API llama_pos llama_kv_cache_seq_pos_max(ctx));
    DEPRECATED(LLAMA_API void llama_kv_cache_defrag(ctx));
    DEPRECATED(LLAMA_API bool llama_kv_cache_can_shift(ctx));
    DEPRECATED(LLAMA_API void llama_kv_cache_update(ctx));

    // works with `ctx.kv_self` - backwards compatible with `master`
    LLAMA_API void llama_kv_self_clear(ctx);
    LLAMA_API bool llama_kv_self_seq_rm(ctx);
    LLAMA_API void llama_kv_self_seq_cp(ctx);
    LLAMA_API void llama_kv_self_seq_keep(ctx);
    LLAMA_API void llama_kv_self_seq_add(ctx);
    LLAMA_API void llama_kv_self_seq_div(ctx);
    LLAMA_API llama_pos llama_kv_self_seq_pos_max(ctx);
    LLAMA_API void llama_kv_self_defrag(ctx);
    LLAMA_API bool llama_kv_self_can_shift(ctx);
    LLAMA_API void llama_kv_self_update(ctx);

    // TODO: llama_kv_cache API
    // can be implemented in a later PR
    // new API to access the KV cache instance
    struct llama_kv_cache;

    LLAMA_API struct llama_kv_cache * llama_get_kv_self(ctx)
    LLAMA_API void                    llama_set_kv_self(ctx, kv);
    // allow to clone, free, save, load the kv cache

@slaren Coming back to your comment from earlier: #11110 (review)

• At some point we should abstract eveything needed to model an architecture to a single class (such that each architecture is a subclass of this class)
• After that, llm_type should probably be removed entirely, and each architecture should have its own enum if needed, with a function to return the type as a string (which by default could be " ")

In the OP I have outlined a possible approach to make the implementation more abstract. I have focused primarily on the abstraction of the KV cache and the llama context.

If I understand correctly your suggestion, the idea is to have the compute graph build functions for each of the arches (e.g. build_llama()), become part of llama_model (e.g. implement derived classes llama_model_llama, llama_model_qwen, etc..), which would effectively eliminate the need for llm_build_context. This way, the llama_context would be able to simply call model->build(), instead of relying on the graph to come from "outside". Do I understand correctly the idea?

I haven't really though enough about this to make specific suggestions, but I think the goal should be to have an interface that can be used to define everything necessary to implement a model architecture. Ideally, to add support for a new architecture, it should only be necessary be to define a new class and create a mapping between the architecture name in the GGUF file and this class. There may of course be more classes in the interface, but there should be a single entry point. So this should include more than just the graph build function, it should also include all the functions to load a model, create a context, and everything else that may be necessary to run a model. This interface would also need to be supported by other interfaces such as the KV cache abstraction, and graph building helper functions that are currently in llm_build_context and the other llm_build_* functions.

To do this, I think it would be better to create an abstract interface that contains everything necessary to define a model architecture. I think that's likely to result in a cleaner and more maintainable codebase than using llama_model as a base class. Instead, llama_model (and other classes like llama_context) should use this interface to implement the functionality in llama.cpp. It may also be convenient to have one or more base classes that implement some of the common functionality that is shared between multiple model architectures, but it should not be strictly necessary to use these base classes.

This is of course a very high level suggestion, it will take a lot of work to define all the details.

Thanks for the suggestions. I'll aim to create the abstract model interface and restructure the implementation so that the llm_build_context is no longer needed and all model-specific code is behind the new abstract interface. Will keep hacking on this PR for a while and try to bring it in a more complete state before merging.

This PR is getting close to completion. Here is an update of the new software architecture:

• llama_context is now a base class implementing a cache-less (i.e. without a KV cache) inference. This means that the compute graphs operate on the current batch and there is no state being preserved across encode()/decode() calls. It also provides basic ggml-related functionality such as initializing backends and preparing output buffers for common things, such as logits, embeddings, etc.
• The llama_context now also implements a "graph building" interface llama_graph_i. The idea is that every model will utilize this interface to create its compute graphs. For example, where a model requires an attention block, it now calls llama_graph_i::build_attn() and delegates the logic to the specific instance of the llama_context. This is the main change that was needed to decouple the KV cache from the graph build logic and enables the implementation of new KV caches and model architectures.
• The base llama_context is used for encoder-only models such as BERT since they do not require a KV cache. But also can work with any decode model as well - it will just be slow because there is no cache.
• The llama_context_kv_self class inherits llama_context and adds a llama_kv_cache instance to support self-attention caching. This class is mainly useful for decoder models. In the future, it will support a "per-sequence" KV cache variation which will be utilized in multi-user use cases, because the regular unified KV cache is not optimal in such scenarios.
• The llama_context_recurrent class is used for recurrent models, such as RWKV and Mamba. It utilizes a llama_kv_cache_recurrent instance, which currently is temporary implemented as a regular llama_kv_cache. In the future, we will implement a recurrent-specific cache that is suitable for these architectures. (cc @compilade)
• Before merging, I need to add an encoder-decoder context for models such as T5. I'm still working out the details, but I am thinking along the lines of composing 2 context - one for the encoder and one for the decoder. This should serve as the basic example for multi-modal support later on.
• The model definitions are now fully contained in llama-model.cpp. This includes hparams and data loading + graph builds. We are still utilizing a "build context" like before, but it is now simplified. In follow-up PRs, this will be improved by moving each model definition in a separate class.
• The encode()/decode() implementations of each llama_context currently have a lot of duplicated code. These should be improved in the future. The main purpose of these calls is to perform the micro-batching according to the attention implementation of the context and in order to improve this, there are likely changes needed to the llama_ubatch, llama_sbatch and llama_batch objects. So to avoid increasing the scope of this PR even further, these will be reworked later on, likely within llama : private llama_batch #11875.
• We can now also implement a MLA-specific llama_context_kv_self_mla for R1 models. It can have a customized llama_kv_cache_mla implementation + extra context state and custom attention implementation. (cc @fairydreaming)

Pinging @MollySophia and @compilade if you could run some tests with this branch to check if the RWKV and Mamba models work correctly.

Any suggestions for improving the code are welcome. Hoping to have this ready for review in the next few days.

I've been quite on/off recently, but hopefully I can have a deeper look into this during the weekend.

The llama_context now also implements a "graph building" interface llama_graph_i. The idea is that every model will utilize this interface to create its compute graphs. For example, where a model requires an attention block, it now calls llama_graph_i::build_attn() and delegates the logic to the specific instance of the llama_context. This is the main change that was needed to decouple the KV cache from the graph build logic and enables the implementation of new KV caches and model architectures.

@ggerganov I see that there is an implicit assumption in build_attn() method that full Q, K and V vectors exist. But the main idea of #11446 is to avoid calculation of full Q K V full vectors. How do you plan to handle this case? Shall I add a separate build_attn_mla() method across the codebase starting from the llama_graph_i interface and going through llama_context up to the new llama_context_kv_self_mla?

Also, since the model definition is now fully contained inside llama-model.cpp, I'm wondering if we can maybe implement some sort of "graph check" in the future. The idea is to export a graph for a given model as a "snapshot", then have a CI check to make sure they are not inadvertently changed. Ideally, this can be done without loading model weight or creating llama_context / KV cache.

Currently, we can debug the cgraph using ggml_graph_dump_dot, but this requires loading the model weight which will be impossible for very large models (>70B)

I'm not sure if the idea is worth exploring, but I can create a dedicated issue to discuss more if needed.

The llama_context now also implements a "graph building" interface llama_graph_i. The idea is that every model will utilize this interface to create its compute graphs. For example, where a model requires an attention block, it now calls llama_graph_i::build_attn() and delegates the logic to the specific instance of the llama_context. This is the main change that was needed to decouple the KV cache from the graph build logic and enables the implementation of new KV caches and model architectures.

@ggerganov I see that there is an implicit assumption in build_attn() method that full Q, K and V vectors exist. But the main idea of #11446 is to avoid calculation of full Q K V full vectors. How do you plan to handle this case? Shall I add a separate build_attn_mla() method across the codebase starting from the llama_graph_i interface and going through llama_context up to the new llama_context_kv_self_mla?

@fairydreaming Yes, you should add a new llama_graph_i::build_attn_mla() in case the current signature of llama_graph_i::build_attn() is not suitable for the MLA use case. The default implementation in llama-graph.cpp would be to return an error of not being implemented. You then don't need to make any changes to llama_context - you only need to implement it in llama_context_kv_self_mla.

Looks good overall. Some points I'm thinking about for my vision PR:

• Having a derived class llama_vision_context : llama_context as you said
• Input image tokens will be obtained via llama_batch_ext, they will be passed to llama_vision_context::input_set which can work with pixel values instead of text token
• Output tensor will be saved to llama_context::embd_tensor ==> need to add this to the base class

Overall yes. The details are not yet clear to me completely - I think once the T5 encoder-decoder use case is implemented we will have a more clear picture and a starting point for multi-modal support.

What I am trying to do is to be able to compose the llama_contexts together. For example, if we look at the Whisper model (because I am familiar with it), it is essentially an Encoder + Decoder. So it should be possible to implement it using the same llama_context_enc_dec that we would use for the T5 models. At the same time, we should also have the option to create the encoder/decoder contexts individually. For example, I would be able to say: create a llama_context_enc from this Whisper model. And this context will be used only for encoding audio to embeddings.

Extending this analogy, a vision model is likely to fit in the same llama_context_enc_dec implementation because it is again an Encoder + Decoder stitched together. The specific input type is not relevant to the context's purpose/implementation.

@ggerganov I think there's still one thing missing. There should be an abstract kv cache interface, llama_kv_cache_i or something like this that caches would implement (and llama_context::get_kv_self() would return this type). I see that you initially planned that llama_kv_cache type would serve as a base type for caches, but as far as I can see this is not implemented (the current llama_kv_cache still contains code specific to caching K/V vectors inside and there's no llama_kv_cache_unified).